import numpy as np
from statsmodels.tsa.holtwinters import ExponentialSmoothing

from service.models import PredictionFact, Month
from service.utils import split_qtty_to_weeks, split_predicted_qtty_to_weeks, split_weeks
from users.models import Specialization, SpecializationType
import pandas as pd
class PredictionService:

    def predict_data(self):
        specs = SpecializationType.objects.all()
        for spec in specs:
            data = PredictionFact.objects.select_related('month').filter(specialization=spec,qtty__isnull=False).all().values('month__year','week','qtty')
            if not data:
                continue
            df = pd.DataFrame(list(data)).rename(columns={'month__year':'year'})
            df= df.sort_values(['year','week'])
            # df= df.iloc[:-3,:]
            df = df.groupby(['year','week'],as_index=False).agg({'qtty':'sum'})
            df['date'] = pd.to_datetime(df['year'].astype(str) + df['week'].astype(str) + '1', format='%Y%W%w')
            data = df.set_index('date')['qtty']
            try:
                model_hw = ExponentialSmoothing(data, seasonal='add', seasonal_periods=52).fit()
            except Exception as e:
                model_hw = ExponentialSmoothing(data).fit()
            results = model_hw.forecast(5).reset_index(drop=True)#.rename(columns={'index':'date'})
            last_date = df.sort_values('date').iloc[-1]['date']

            future_date = [last_date + pd.DateOffset(weeks=i) for i in range(1,6)]
            future_df = pd.DataFrame(future_date, columns=['date'])
            # Вычисление года и номера недели
            future_df['year'] = future_df['date'].dt.year
            future_df['month'] = future_df['date'].dt.month
            future_df['week'] = future_df['date'].dt.isocalendar().week
            res = pd.DataFrame()
            res['year'] = future_df['year']
            res['week'] = future_df['week']
            res['month'] = future_df['month']
            res['qtty'] =results
            res=res.astype(int)
            expanded_df = split_weeks(res[['year','week']])
            res = res.merge(expanded_df, on=['year', 'week'])
            res = res.rename(columns={'month_y': 'month'}).drop(columns='month_x')
            res = split_predicted_qtty_to_weeks(res)
            for ind, row in res.iterrows():
                month, created = Month.objects.get_or_create(year=int(row['year']), month=int(row['month']))
                pred, created = PredictionFact.objects.get_or_create(week=int(row['week']), month=month, specialization=spec)
                pred.qtty_predicted = row['qtty']
                pred.save()